# -*- coding: utf-8 -*-
"""
File: udp_random_flood.py
功能：高并发 UDP 随机字节码洪流压测工具（支持随机 MSG_TYPE）
描述：模拟大量随机 UDP 数据包，MSG_TYPE 在 1 和 3 之间随机选择，用于全面测试服务端协议处理能力
特点：
  - 每个数据包内容为随机字节（无规律）
  - MSG_TYPE 随机为 1 或 3
  - 可配置包大小、发送频率、并发线程
  - 支持带协议头或纯随机流量
"""

import struct
import socket
import time
import random
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

# ==================== 配置参数 ====================
SERVER_IP = "120.26.136.5"
SERVER_PORT = 50001
TIMEOUT = 5               # UDP 接收超时（秒）
NUM_THREADS = 10          # 线程池大小
PACKETS_PER_BATCH = 100   # 每批次发送的数据包数量
SEND_INTERVAL = 1         # 发送间隔（秒）

# ✅ 随机字节包大小范围
MIN_PACKET_SIZE = 32
MAX_PACKET_SIZE = 1024

# ✅ MSG_TYPE 可选值
MSG_TYPE_OPTIONS = [1, 3]

# 总运行时间：24 小时
RUN_DURATION = 24 * 60 * 60  # 秒

# UDP 协议头控制
USE_HEADER = True           # 是否添加协议头
START_CODE = 0xFFFFFFA1
MSG_HEADER_FORMAT = ">III"  # StartCode, MsgType, MsgLen (big-endian)


# ==================== 工具函数 ====================
def generate_random_bytes():
    """生成指定范围内的随机字节"""
    size = random.randint(MIN_PACKET_SIZE, MAX_PACKET_SIZE)
    return bytes(random.getrandbits(8) for _ in range(size))


def create_packet():
    """构建带协议头的 UDP 数据包，MSG_TYPE 随机选择"""
    payload = generate_random_bytes()

    if USE_HEADER:
        msg_type = random.choice(MSG_TYPE_OPTIONS)  # ✅ 随机选择 1 或 3
        msg_len = len(payload)
        header = struct.pack(MSG_HEADER_FORMAT, START_CODE, msg_type, msg_len)
        return header + payload
    else:
        return payload


def send_packet(data):
    """发送单个 UDP 包并尝试接收响应（可选）"""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.settimeout(TIMEOUT)
            sock.sendto(data, (SERVER_IP, SERVER_PORT))
            # 可选：接收响应（用于判断通断）
            # response, _ = sock.recvfrom(65535)
            return True
    except socket.timeout:
        return False
    except Exception as e:
        # print(f"Send error: {e}")
        return False


def send_batch():
    """发送一批随机 UDP 数据包"""
    packets = [create_packet() for _ in range(PACKETS_PER_BATCH)]
    successful = 0
    total = len(packets)

    start_time = time.time()

    with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:
        future_to_packet = {executor.submit(send_packet, packet): packet for packet in packets}
        for future in as_completed(future_to_packet):
            if future.result():
                successful += 1

    elapsed = time.time() - start_time
    qps = total / elapsed if elapsed > 0 else 0
    success_rate = (successful / total) * 100

    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{now}] Sent {total}, Success: {successful}, "
          f"QPS: {qps:.2f}, Success Rate: {success_rate:.1f}%, "
          f"Time: {elapsed:.3f}s")

    return successful, total


# ==================== 主逻辑 ====================
def main():
    print(f"🚀 UDP 随机字节洪流压测开始...")
    print(f"目标：每 {SEND_INTERVAL} 秒发送 {PACKETS_PER_BATCH} 个随机 UDP 包")
    print(f"包大小：{MIN_PACKET_SIZE} ~ {MAX_PACKET_SIZE} 字节")
    print(f"MSG_TYPE: 随机从 {MSG_TYPE_OPTIONS} 中选择")
    print(f"服务器: {SERVER_IP}:{SERVER_PORT}, 线程数: {NUM_THREADS}")
    print(f"协议头: {'启用' if USE_HEADER else '禁用'}")
    print(f"将持续运行 {RUN_DURATION // 3600} 小时...")

    start_time = time.time()
    total_success = 0
    total_packets = 0

    while time.time() - start_time < RUN_DURATION:
        success, count = send_batch()
        total_success += success
        total_packets += count

        # 精确对齐发送时间
        next_send_time = (start_time + (int((time.time() - start_time) / SEND_INTERVAL) + 1) * SEND_INTERVAL)
        sleep_time = max(0, next_send_time - time.time())
        time.sleep(sleep_time)

    # 最终统计
    run_time = time.time() - start_time
    overall_qps = total_packets / run_time if run_time > 0 else 0
    overall_success_rate = (total_success / total_packets * 100) if total_packets > 0 else 0

    print("\n" + "="*60)
    print("✅ 压测完成！")
    print(f"总运行时间: {run_time//3600:.0f}h {(run_time%3600)//60:.0f}m")
    print(f"发送间隔: {SEND_INTERVAL}s")
    print(f"每批包数: {PACKETS_PER_BATCH}")
    print(f"总批次数: {total_packets // PACKETS_PER_BATCH}")
    print(f"总发送包数: {total_packets:,}")
    print(f"成功响应数: {total_success:,}")
    print(f"整体成功率: {overall_success_rate:.2f}%")
    print(f"平均 QPS: {overall_qps:.2f}")
    print("="*60)


if __name__ == "__main__":
    main()